# -*- coding: utf-8 -*-
'''
Created on May 20, 2012

@author: tel
'''
import sys, os, time, numpy, site, re
#sys.path.append('/home/tel/workspace/pmcce/src')
site.addsitedir('../')
from collections import namedtuple

from sqlalchemy import create_engine, event, MetaData
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool

import makedb.levels as levels

def _fk_pragma_on_connect(dbapi_con, con_record):
    dbapi_con.execute('pragma foreign_keys=ON')

def LoadSession(db):
    Session = sessionmaker(bind=db)
    session = Session()
    return session
    
def Nearby_Mutants(session):
    for prot, res in session.query(levels.ProtData, levels.ResData).filter(levels.ResData.seqadv!='').yield_per(10):
        near_res = {}
        for atom in res.Atoms:
            if atom.atom=='CB':
                for near_atom in prot.Atoms:
                    if atom.Dist(near_atom) < 10 and near_atom.atom=='CB':
                        near_res[str(near_atom.Res.res)+str(near_atom.Res.chainid)+str(near_atom.Res.resn)] = near_atom.Res
        print 'this mutant:'
        print '%s/e%s\t%s\t%s\t%s\t' % (prot.pdb,res.Run.EPSILON_PROT,res.res,res.chainid,res.resn)
        print 'is within 10 A of these other residues:'
        for r in near_res.values():
            print '%s/e%s\t%s\t%s\t%s\t' % (prot.pdb,res.Run.EPSILON_PROT,r.res,r.chainid,r.resn)

#for prot,run,res in session.query(levels.ProtData, levels.RunData, levels.ResData).join(levels.ProtData.Runs).join(levels.RunData.Ress).filter(levels.ResData.charge!=None).yield_per(10): 
def SumCrgAt(session, titr):
    '''
    gives the ionization of every residue at ph/em = titr
    '''
    print levels.ResData.StrHeaders(), '\tSum_Crg_at_%.3f' % titr
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        print res, '\t%.3f' % res.SumCrgAt(titr)
        
def DPKGreaterThan(session, dpk):
    '''
    gives every residue for which the pKa differs from the solution value by dpk or more, and the pKa
    '''
    print levels.ResData.StrHeaders(), '\tpKa\tDpKa'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        if abs(res.dpk)>dpk:
            print res, '\t%.3f\t%.3f' % (res.__getattribute__('pKa/Em'),res.dpk)
            
def DsolGreaterThan(session, dsol):
    '''
    gives every residue for which the pKa differs from the solution value by dpk or more, and the pKa
    '''
    print levels.ResData.StrHeaders(), '\tDsol'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        if abs(res.dsol)>dsol:
            print res, '\t%.3f' % res.dsol

def RMSDGreaterThan(session, rmsd):
    print levels.ResData.StrHeaders(), '\tRMSD'
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        titrs = sorted(res.Confs[1].occupancy.keys())
        initial = res.MostOccAt(titrs[0])
        final = res.MostOccAt(titrs[-1])
        RMSD = initial.RMSD(final,session)
        if RMSD>rmsd:
            print res, '\t%.3f' % RMSD
            
def CrystalConformers(session):
    '''
    finds confomers for each residue that have heavy atom positions identical to the input (crystal) structure
    prints out conformers by number
    '''
    print levels.ResData.StrHeaders(), '\tConformers_with_Crystal_Heavy_Atoms'
    for res in session.query(levels.ResData).yield_per(10):
        crysconfs = [1]
        for i,conf in enumerate(res.Confs[2:]):
            if res.Confs[1].RMSD(conf,session)==0:
                crysconfs.append(i+2)
        print res, '\t%s' % crysconfs

def SumCrgNonmonotonic(session, tolerance=.05):
    '''
    gives every residue for which the ionization goes up and down across the titation, and the ionization at all titration points
    will not report as nonmonotonic if variation is within the tolerance
    '''
    print levels.ResData.StrHeaders()
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        sumcrg = sorted(res.sumcrg.items()) #sumcrg is stored as a dict
        ionization = zip(*sumcrg)[1]
        deltas = map(lambda x,y: y-x, ionization[:-1], ionization[1:])
        big_pos = False
        big_neg = False
        for delta in deltas:
            if delta>tolerance:
                if big_neg==True:
                    print res
                    print ionization
                    break
                else:
                    big_pos=True
            elif delta<-tolerance:
                if big_pos==True:
                    print res
                    print ionization
                    break
                else:
                    big_neg=True
                    
def CmpMFE(session):
    '''
    function to test MFE functionality in mccedb vs MFE in vanilla
    '''
    print levels.ResData.MFE.StrHeaders()
    for res in session.query(levels.ResData).filter(levels.ResData.charge!=None).yield_per(10):
        res.GetMFE()
        cmpline = ''
        for attr in res.MFE.Ionization.mfe_attrs:
            cmpline += '%.3f\t' % res.__getattribute__(attr)
        print res
        print cmpline
        print res.mfe
        
def Jiggle(session):
    '''
    prints jiggle for every residue
    '''
    print levels.ResData.StrHeaders(), '\tJiggle'
    for res in session.query(levels.ResData).yield_per(10):
        if res.__getattribute__('pKa/Em')!=None:
            print res, '\t%s' % res.GetJiggle(atpk=True), '\t%s' % res.dsol, '\t%s' % (res.e_r + res.e_b)
                    
def AtomInfo(session, res, resn, chainid, atomn):
    atom = session.query(levels.AtomData).join(levels.ResData).filter(levels.ResData.res==res).filter(levels.ResData.resn==resn).filter(levels.ResData.chainid==chainid).filter(levels.AtomData.atomn==atomn).first()
    print atom.StrHeaders(), 'X\tY\tZ'
    print atom, ('%.3f\t%.3f\t%.3f\t' % (atom.x, atom.y, atom.z))
            
if __name__=='__main__':
    funcdict = {'SumCrgAt':SumCrgAt,'DPKGreaterThan':DPKGreaterThan,'DsolGreaterThan':DsolGreaterThan,'RMSDGreaterThan':RMSDGreaterThan,'CrystalConformers':CrystalConformers,'SumCrgNonmonotonic':SumCrgNonmonotonic,'AtomInfo':AtomInfo,'CmpMFE':CmpMFE,'Jiggle':Jiggle}
    #this next part recreates all of the string keys in funcdict so that they are completely in lowercase
    tmpdict = {}
    for key, value in funcdict.iteritems():
        tmpdict[key.lower()] = value
    funcdict = tmpdict
    try:
        dbpath = sys.argv[1]
        func = funcdict[sys.argv[2].lower()]
    except (KeyError, IndexError):
        print 'Usage:'
        print 'path-to-database function [argument]\n'
        print 'Possible functions:'
        for key in funcdict.keys():
            print key
        sys.exit(0)
    arg = None
    if len(sys.argv) > 3:
        arg = sys.argv[3]
    dir = '/home/tel/test/db1'
    db = create_engine('sqlite:///' + dbpath, poolclass=NullPool)
    event.listen(db, 'connect', _fk_pragma_on_connect)
    meta= MetaData(db)
    
    for level in levels.hierarchy:
        level.Table(meta=meta, db=db, reloadt=True)
    session = LoadSession(db)
    
    #Nearby_Mutants(session)
    if arg!=None:
        try:
            func(session, float(arg))
        except TypeError:
            func(session)
    else:
        func(session)
